import torch
import torch.nn as nn
from .utils import conv, deconv
from compressai.layers import GSDN
from compressai.ops import quantize_ste
from .base import (
    CompressionModel,
)

class Ali2021(CompressionModel):
    def __init__(self, N, M, **kwargs):
        super().__init__(**kwargs)
        self.g_a = nn.Sequential(
            conv(3, N),
            GSDN(N),
            conv(N, N),
            GSDN(N),
            conv(N, N),
            GSDN(N),
            conv(N, M),
        )

        self.g_s = nn.Sequential(
            deconv(M, N),
            GSDN(N, inverse=True),
            deconv(N, N),
            GSDN(N, inverse=True),
            deconv(N, N),
            GSDN(N, inverse=True),
            deconv(N, 3),
        )

        self.h_a = nn.Sequential(
            conv(M, N, stride=1, kernel_size=3),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            conv(N, N),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            conv(N, N),
        )

        self.h_s = nn.Sequential(
            deconv(N, N),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            deconv(N, N),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            conv(N, N, stride=1, kernel_size=3),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            conv(N, M * 2, stride=1, kernel_size=1),
        )

    def forward(self, x):
        y = self.g_a(x)
        y_tilde = y + torch.empty_like(y).uniform_(-float(0.5), float(0.5))
        y_tilde2 = quantize_ste(y)

        z = self.h_a(y)
        z_tilde = z + torch.empty_like(z).uniform_(-float(0.5), float(0.5))
        z_hat = self.h_s(z_tilde)

        x_tilde = self.g_s(y_tilde2).clamp(0., 1.)